#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <openssl/evp.h>
#include <unistd.h>
#include "openssl/ssl.h"
#include "openssl/err.h"
#include "openssl/sm4.h"
#include "openssl/engine.h"

// 参考自：http://gmssl.org/docs/evp-api.html

#define BENCHMARK_TIME 10000

enum MODE {
	ECB = 0,
    CBC = 1,
};

int encrypt(unsigned char *plaintext, int plaintext_len, unsigned char *key,
unsigned char *iv, unsigned char *ciphertext, enum MODE mode)
{
    EVP_CIPHER_CTX *ctx;
    int len;
    int ciphertext_len;
    const EVP_CIPHER *cipher = NULL;

	if (mode == ECB) {
		iv = NULL;
		cipher = EVP_sm4_ecb();
	}
	else if (mode == CBC) {
		cipher = EVP_sm4_cbc();
	}
	else {
		printf("Doesn't support this mode :[%d].\n", mode);
	}

    if(!(ctx = EVP_CIPHER_CTX_new())) {
        printf("EVP initailized failed.\n");
        return -1;
    }

    if(1 != EVP_EncryptInit_ex(ctx, cipher, NULL, key, iv)) {
        printf("EVP Decrypt Init failed.\n");
        return -1;
    }

    if(1 != EVP_EncryptUpdate(ctx, ciphertext, &len, plaintext, plaintext_len)) {
        printf("EVP Decrypt update failed.\n");
        return -1;
    }
    ciphertext_len = len;

    if(1 != EVP_EncryptFinal_ex(ctx, ciphertext + len, &len)) {
        printf("EVP Decrypt final failed\n");
        return -1;
    }
    ciphertext_len += len;

    /* Clean up */
    EVP_CIPHER_CTX_free(ctx);

    return ciphertext_len;
}

int decrypt(unsigned char *ciphertext, int ciphertext_len, unsigned char *key,
unsigned char *iv, unsigned char *plaintext, enum MODE mode)
{
    EVP_CIPHER_CTX *ctx;
    int len;
    int plaintext_len;
    const EVP_CIPHER *cipher = NULL;

	if (mode == ECB) {
		iv = NULL;
		cipher = EVP_sm4_ecb();
	}
	else if (mode == CBC) {
		cipher = EVP_sm4_cbc();
	}
	else {
		printf("Doesn't support this mode :[%d].\n", mode);
	}

    /* Create and initialise the context */
    if(!(ctx = EVP_CIPHER_CTX_new())) {
        printf("EVP initailized failed.\n");
        return -1;
    }

    if(1 != EVP_DecryptInit_ex(ctx, cipher, NULL, key, iv)) {
        printf("EVP Decrypt Init failed.\n");
        return -1;
    }

    if(1 != EVP_DecryptUpdate(ctx, plaintext, &len, ciphertext, ciphertext_len)) {
        printf("EVP Decrypt update failed.\n");
        return -1;
    }
    plaintext_len = len;

    if(1 != EVP_DecryptFinal_ex(ctx, plaintext + len, &len)) {
        printf("EVP Decrypt final failed\n");
        return -1;
    }
    plaintext_len += len;

    /* Clean up */
    EVP_CIPHER_CTX_free(ctx);

    return plaintext_len;
}

char *randstr(char *str, const int len) {
    srand(time(NULL));
    int i;
    for (i = 0; i < len; ++i)
    {
        switch ((rand() % 3))
        {
        case 1:
            str[i] = 'A' + rand() % 26;
            break;
        case 2:
            str[i] = 'a' + rand() % 26;
            break;
        default:
            str[i] = '0' + rand() % 10;
            break;
        }
    }
    str[++i] = '\0';
    return str;
}

void benchsm4_enc_dec(int size, enum MODE mode) {
	// 测试用的秘钥和初始向量
    unsigned char key[] = "1234567890123456";
    unsigned char iv[]  = "0000000000000000";

    int cryptsize = 0;
    if (size < 1024)
        cryptsize = 1024;
    else if (size < 8000)
        cryptsize = 4096;
    else 
        cryptsize = 20000;

	char m[4] = {0};
	if (mode == CBC) strcpy(m, "ECB");
	else if (mode == ECB) strcpy(m, "CBC");

    char *msg = (char *)malloc(size);
    unsigned char *plaintext = (char *)malloc(cryptsize);
    unsigned char *decrypttext = (char *)malloc(cryptsize);
	memset(msg, 0, size);
	memset(plaintext, 0, cryptsize);
	memset(decrypttext, 0, cryptsize);
	randstr(msg, size-1);
//	printf("msg :[%s]\n", msg);
    int msg_len = strlen(msg);
	int decrypt_len = 0;	

	clock_t start, finish;
    double totaltime;
	
	start = clock();
	for (int i = 0; i < BENCHMARK_TIME; ++i) {
	    decrypt_len = encrypt(msg, msg_len, key, iv, decrypttext, mode);
	}
	finish = clock();
	totaltime = (double)(finish - start) / CLOCKS_PER_SEC / BENCHMARK_TIME;
    printf("%s 加密[%d]Bytes时间测试（测试次数：%d）: %f\n", m, size-1, BENCHMARK_TIME, totaltime);

    /*printf("Decrypt result: [");
    for (int i = 0; i < decrypt_len; ++i) {
         printf("%02x", decrypttext[i]);
    }
    printf("]\n");*/

	int plaintext_len = 0;
	start = clock();
	for (int i = 0; i < BENCHMARK_TIME; ++i) {
    	plaintext_len = decrypt(decrypttext, decrypt_len, key, iv, plaintext, mode);
	}
	finish = clock();
	totaltime = (double)(finish - start) / CLOCKS_PER_SEC / BENCHMARK_TIME;
    printf("%s 解密[%d]Bytes时间测试（测试次数：%d）: %f\n", m,size-1, BENCHMARK_TIME, totaltime);

    //printf("%d %s\n", plaintext_len, plaintext);
	if (msg) free(msg);
	if (plaintext) free(plaintext);
	if (decrypttext) free(decrypttext);
}

int main(int argc, char **argv)
{
	benchsm4_enc_dec(9, ECB);  
	benchsm4_enc_dec(1025, ECB);  
	benchsm4_enc_dec(8193, ECB);  
	benchsm4_enc_dec(9, CBC);  
	benchsm4_enc_dec(1025, CBC);  
	benchsm4_enc_dec(8193, CBC);  
    return 0;
}
